import re
import time

import pymysql
from bs4 import BeautifulSoup
from flask import Flask, jsonify, request, render_template, flash, session, json
from pymysql import OperationalError
from pymysql.cursors import DictCursor
from flask_cors import CORS
from selenium import webdriver
from selenium.common import TimeoutException
from selenium.webdriver.common.by import By
from selenium.webdriver.support.wait import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC

app = Flask(__name__)
CORS(app)
app.config['SECRET_KEY'] = '#8asdftt7'

def get_db_connection():
    connection = pymysql.connect(
        host='localhost',
        user='root',
        password='123456',
        database='flask',
        cursorclass=DictCursor
    )
    return connection


# 首页
@app.route('/')
def index():
    return render_template('login.html')

@app.route('/next')
def next():
    return render_template('stock_analysis_fixed.html')

@app.route('/stock')
def menu():
    return render_template('stock.html')

@app.route('/shenzhen')
def shenzhen():
    return render_template('stock_analysis_shenzhen.html')

@app.route('/tospider')
def tospider():
    return render_template('spider.html')

#登录
@app.route('/login',methods=['POST'])
def login():
    connection = get_db_connection()
    with connection.cursor() as cursor:
        cursor.execute('SELECT * FROM user')
        users = cursor.fetchall()
        for i in users:
            if i['username'] == request.form.get('username') and i['password'] == request.form.get('password'):
                connection.close()
                return render_template('test.html')
        connection.close()
        flash('用户名或密码错误')
        return render_template('login.html')

#注册
@app.route('/register',methods=['POST'])
def register():
    connection = get_db_connection()
    with connection.cursor() as cursor:
        cursor.execute('select username from user')
        users = cursor.fetchall()
        username = request.form.get('username')
        password = request.form.get('password')
        for i in users:
            if i['username'] == username:
                flash("用户已经存在，请更换用户名")
                return render_template('register.html')
        while (len(username.replace(" ", "")) == 0 or len(password.replace(" ", "")) == 0):
            flash("请输入非空账号和密码")
            return render_template('register.html')
        cursor.execute('insert into user (username, password) values (%s, %s)',(username, password))
        connection.commit()
        flash("注册成功，请登录")
        return render_template('login.html')


# 获取所有表名
@app.route('/tables', methods=['GET'])
def get_tables():
    connection = get_db_connection()
    try:
        with connection.cursor() as cursor:
            cursor.execute("SHOW TABLES")
            tables = [row[f'Tables_in_{connection.db.decode()}'] for row in cursor.fetchall()]
        return jsonify(tables)
    finally:
        connection.close()


# 获取表结构
@app.route('/<table>/structure', methods=['GET'])
def get_table_structure(table):
    connection = get_db_connection()
    try:
        with connection.cursor() as cursor:
            cursor.execute(f"DESCRIBE {table}")
            structure = cursor.fetchall()
        return jsonify(structure)
    finally:
        connection.close()


# 通用分页查询（含搜索、排序）
@app.route('/<table>/find', methods=['POST'])
def find_records(table):
    connection = get_db_connection()
    try:
        data = request.get_json() or {}
        current_page = int(request.args.get('currentPage', 1))
        page_size = int(request.args.get('pageSize', 10))
        begin = (current_page - 1) * page_size

        # 获取表结构（验证字段存在性）
        with connection.cursor() as cursor:
            cursor.execute(f"DESCRIBE {table}")
            columns = [col['Field'] for col in cursor.fetchall()]

        # 构建搜索条件
        conditions = []
        params = []
        search_value = data.get('search')
        if search_value:
            or_conditions = [f"{col} LIKE %s" for col in columns]
            conditions.append("(" + " OR ".join(or_conditions) + ")")
            params.extend([f"%{search_value}%"] * len(columns))

        # 处理额外过滤条件（前端传递的其他键值对）
        for key, value in data.items():
            if key != 'search' and value:
                if key in columns:
                    conditions.append(f"{key} LIKE %s")
                    params.append(f"%{value}%")

        where_clause = "WHERE " + " AND ".join(conditions) if conditions else ""

        # 处理排序（新增功能）
        sort_field = request.args.get('sortField', '').strip()
        sort_order = request.args.get('sortOrder', 'asc').lower()
        # 验证排序字段合法性
        if sort_field and sort_field not in columns:
            return jsonify({
                'success': False,
                'message': f'排序字段 "{sort_field}" 不存在'
            }), 400
        order_by = f"ORDER BY `{sort_field}` {sort_order}" if sort_field else ""

        # 数据查询
        query = f"SELECT * FROM {table} {where_clause} {order_by} LIMIT %s, %s"
        params.extend([begin, page_size])

        with connection.cursor() as cursor:
            cursor.execute(query, params)
            rows = cursor.fetchall()

            # 总记录数查询（包含排序条件，确保计数准确）
            count_query = f"SELECT COUNT(*) as count FROM {table} {where_clause} {order_by}"
            cursor.execute(count_query, params[:-2])  # 排除LIMIT参数
            total_count = cursor.fetchone()['count']

        return jsonify({
            'rows': rows,
            'totalCount': total_count
        })

    except OperationalError as e:
        print(f"数据库操作错误: {e}")
        connection.rollback()
        return jsonify({
            'success': False,
            'message': '数据库操作失败',
            'error': str(e)
        }), 500

    except Exception as e:
        print(f"系统错误: {e}")
        connection.rollback()
        return jsonify({
            'success': False,
            'message': '系统处理失败',
            'error': str(e)
        }), 500

    finally:
        connection.close()


# 通用添加记录
@app.route('/<table>/add', methods=['POST'])
def add_record(table):
    connection = get_db_connection()
    try:
        data = request.get_json()
        if not data:
            return jsonify({'success': False, 'message': '数据不能为空'}), 400

        # 获取表结构，验证字段
        with connection.cursor() as cursor:
            cursor.execute(f"DESCRIBE {table}")
            columns = [col['Field'] for col in cursor.fetchall()]

        # 过滤不存在的字段（特别注意：id字段不应出现在添加数据中）
        valid_data = {k: v for k, v in data.items() if k in columns and k != 'id'}
        if not valid_data:
            return jsonify({'success': False, 'message': '没有有效字段'}), 400

        # 构建SQL
        fields = ', '.join(valid_data.keys())
        placeholders = ', '.join(['%s'] * len(valid_data))
        values = list(valid_data.values())

        query = f"INSERT INTO {table} ({fields}) VALUES ({placeholders})"

        with connection.cursor() as cursor:
            cursor.execute(query, values)
            connection.commit()

        return jsonify({'success': True, 'message': '添加成功'})
    except Exception as e:
        connection.rollback()
        return jsonify({'success': False, 'message': str(e)}), 500
    finally:
        connection.close()


# 通用更新记录
@app.route('/<table>/update', methods=['POST'])
def update_record(table):
    connection = get_db_connection()
    try:
        record_id = request.args.get('id')
        data = request.get_json()
        if not record_id or not data:
            return jsonify({'success': False, 'message': '缺少参数'}), 400

        # 获取表结构
        with connection.cursor() as cursor:
            cursor.execute(f"DESCRIBE {table}")
            columns = [col['Field'] for col in cursor.fetchall()]

            # 检查记录是否存在
            cursor.execute(f"SELECT * FROM {table} WHERE id = %s", (record_id,))
            record = cursor.fetchone()
            if not record:
                return jsonify({'success': False, 'message': '记录不存在'}), 404

        # 过滤不存在的字段并构建SET子句
        set_clauses = []
        values = []
        for key, value in data.items():
            if key in columns and key != 'id':  # 不更新ID
                set_clauses.append(f"{key} = %s")
                values.append(value)

        if not set_clauses:
            return jsonify({'success': False, 'message': '没有可更新的字段'}), 400

        set_clause = ', '.join(set_clauses)
        values.append(record_id)  # 最后一个参数是ID

        query = f"UPDATE {table} SET {set_clause} WHERE id = %s"

        with connection.cursor() as cursor:
            cursor.execute(query, values)
            connection.commit()

        return jsonify({'success': True, 'message': '更新成功'})
    except Exception as e:
        connection.rollback()
        return jsonify({'success': False, 'message': str(e)}), 500
    finally:
        connection.close()


@app.route('/<table>/delete', methods=['POST', 'DELETE'])  # 支持两种方法
def delete_record(table):
    connection = get_db_connection()
    try:
        if request.method == 'DELETE':  # 处理单个删除
            record_id = request.args.get('id')
            if not record_id:
                return jsonify({'success': False, 'message': '缺少ID参数'}), 400
            ids = [record_id]
        else:  # 处理批量删除
            ids = request.get_json()  # 直接获取ID数组
            if not ids or not isinstance(ids, list):
                return jsonify({'success': False, 'message': '无效的ID列表'}), 400

        with connection.cursor() as cursor:
            # 检查记录是否存在（可选）
            placeholders = ', '.join(['%s'] * len(ids))
            cursor.execute(f"DELETE FROM {table} WHERE id IN ({placeholders})", ids)
            deleted_count = cursor.rowcount
            connection.commit()

        return jsonify({'success': True, 'message': f'成功删除 {deleted_count} 条记录'})
    except Exception as e:
        connection.rollback()
        return jsonify({'success': False, 'message': str(e)}), 500
    finally:
        connection.close()



#爬虫
def find_next_page_button(driver):
    """智能查找下一页按钮，支持多种常见模式"""
    strategies = [
        {
            'name': '东方财富网模式',
            'xpath': '//*[@id="mainc"]/div/div/div[4]/div/a[last()]',
            'text_check': lambda text: text == '>'
        },
        {
            'name': '文本"下一页"按钮',
            'xpath': '//a[contains(text(), "下一页") or contains(text(), "Next")]',
            'text_check': lambda text: True
        },
        {
            'name': '分页容器中的下一页',
            'xpath': '//div[contains(@class, "pagination") or contains(@class, "page")]//a[last()]',
            'text_check': lambda text: '下一页' in text or 'next' in text.lower() or text == '>'
        },
        {
            'name': 'aria-label下一页按钮',
            'xpath': '//a[@aria-label="下一页"]',
            'text_check': lambda text: True
        }
    ]

    for strategy in strategies:
        try:
            buttons = driver.find_elements(By.XPATH, strategy['xpath'])
            for button in buttons:
                if button.is_displayed() and button.is_enabled():
                    button_text = button.text.strip()
                    if strategy['text_check'](button_text):
                        return button, strategy['name']
        except Exception as e:
            print(f"尝试策略 '{strategy['name']}' 失败: {e}")

    return None, None


def analyze_page_for_pagination(driver):
    """分析页面结构，识别分页元素"""
    print("正在分析页面结构以识别分页元素...")

    pagination_containers = driver.find_elements(By.XPATH,
                                                 '//div[contains(@class, "pagination") or contains(@class, "page")]')

    if not pagination_containers:
        pagination_containers = driver.find_elements(By.XPATH,
                                                     '//div[contains(@id, "page") or contains(@class, "pager")]')

    if not pagination_containers:
        print("未识别到分页容器")
        return None, None

    print(f"识别到 {len(pagination_containers)} 个可能的分页容器")

    for container in pagination_containers:
        container_html = container.get_attribute('innerHTML')

        if "下一页" in container_html or "Next" in container_html or ">" in container_html:
            try:
                next_buttons = container.find_elements(By.XPATH,
                                                       './/a[contains(text(), "下一页") or contains(text(), "Next") or text()=">"]')
                if next_buttons:
                    return next_buttons[0], "基于文本识别"

                clickable_elements = container.find_elements(By.XPATH, './/a')
                if len(clickable_elements) > 1:
                    return clickable_elements[-1], "基于位置识别"
            except Exception as e:
                print(f"分析分页容器时出错: {e}")

    return None, None


def wait_for_element(driver, xpath, timeout=10):
    try:
        return WebDriverWait(driver, timeout).until(
            EC.presence_of_element_located((By.XPATH, xpath))
        )
    except TimeoutException:
        print(f"等待元素超时: {xpath}")
        return None


@app.route('/spider', methods=['POST'])
def spider():
    try:
        data = request.json
        url = data.get('url', "https://quote.eastmoney.com/center/gridlist.html#hs_a_board")
        filename = data.get('filename', 'spider_result.csv')
        file_type = data.get('fileType', 'csv')
        encoding = data.get('encoding', 'utf-8')

        all_data = []
        total_pages = 0  # 初始化变量，避免未定义错误

        # 配置驱动选项
        options = webdriver.EdgeOptions()
        options.add_argument('--start-maximized')
        options.add_argument('--disable-notifications')

        driver = webdriver.Edge(options=options)

        try:
            driver.get(url)

            # 等待页面加载
            time.sleep(3)

            def collect_data(html_content):
                soup = BeautifulSoup(html_content, 'lxml')
                trs = soup.find_all('tr')
                page_data = []
                for tr in trs[1:]:
                    tds = tr.find_all(['th', 'td'])
                    row = [td.get_text(strip=True) for td in tds]
                    if row:
                        page_data.append(row)
                return page_data[1:]

            num = 1
            while True:
                try:
                    print(f"正在处理第 {num} 页...")

                    # 等待表格加载
                    wait_for_element(driver, '//table')
                    html = driver.page_source

                    page_data = collect_data(html)
                    all_data.extend(page_data)
                    total_pages = num  # 更新总页数

                    print(f"第 {num} 页收集成功，共 {len(page_data)} 条记录")

                    # 查找下一页按钮
                    next_button, strategy_name = find_next_page_button(driver)

                    if not next_button:
                        next_button, strategy_name = analyze_page_for_pagination(driver)

                    if next_button:
                        if next_button.text != ">":
                            print("数据抓取完毕！")
                            break
                        print(f"使用策略 '{strategy_name}' 找到下一页按钮")

                        # 滚动到按钮位置
                        driver.execute_script("arguments[0].scrollIntoView();", next_button)
                        time.sleep(1)

                        try:
                            next_button.click()
                        except Exception as e:
                            print(f"直接点击失败，尝试JavaScript点击: {e}")
                            driver.execute_script("arguments[0].click();", next_button)

                        print(f"已点击下一页，等待页面加载...")
                        time.sleep(2)
                        num += 1
                    else:
                        print("未找到下一页按钮，爬取完成")
                        break

                except Exception as e:
                    print(f"处理单页时出错: {e}")
                    break

        finally:
            # 确保浏览器总是会被关闭
            driver.quit()

        # 处理数据格式
        if file_type == 'json':
            result = json.dumps(all_data, ensure_ascii=False, indent=2)
        elif file_type == 'csv':
            csv_data = []
            for row in all_data:
                csv_row = ','.join([f'"{cell}"' if ',' in str(cell) else str(cell) for cell in row])
                csv_data.append(csv_row)
            result = '\n'.join(csv_data)
        else:
            result = '\n'.join(['\t'.join(map(str, row)) for row in all_data])

        return jsonify({
            'success': True,
            'data': result,
            'totalPages': total_pages,  # 使用小写的 total_pages
            'message': f'成功爬取{total_pages}页数据'
        })

    except Exception as e:
        print(f"爬取过程出错: {e}")
        return jsonify({
            'success': False,
            'message': f'爬取失败: {str(e)}',
            'totalPages': 0  # 确保返回的 JSON 中包含 totalPages
        })

@app.route('/redirect_login')
def redirect_login():
    return render_template('login.html')

@app.route('/redirect_register')
def redirect_register():
    return render_template('register.html')

if __name__ == '__main__':
    app.run(port=8080, debug=True)